#!/usr/bin/python
"""
This program calculates tetranucleotide frequencies from a given multi-fasta
file and reports the z score

z = (x - mu) / rho

Usage: dissertation_TetraNTCalculator.py multi.fasta tetra.list

tetra.list file should contain a list of tetranucleotide combinations like this:
AAAA
AAAC
AAAG
AAAT
AACA
AACC
AACG
AACT
AAGA
AAGC
.
.
.
TTCG
TTCT
TTGA
TTGC
TTGG
TTGT
TTTA
TTTC
TTTG
TTTT
And should contain a total of 256 combinations.

Note: currently it prints z score in a tab-delimited format like this:
EM7JFSU01D21YQ	-0.471361238918	-0.471361238918	-0.471361238918	3.91658338519
The idea is to use this script to bin metagenomic reads by tetra-nt freq among
other components such as G+C% and other things. I attempt to use z score because
it is a normalized score instead of a raw score which can change based on length
of the sequence.

"""

import sys
import numpy
from Bio import SeqIO
from Bio import Motif
from Bio.Alphabet import IUPAC
from Bio.Seq import Seq

def zscore(x, u, r):
    z = (x - u) / float(r)
    return z

genome = SeqIO.parse(sys.argv[1], "fasta")

tetrafile = open(sys.argv[2], "rU")
tetras = tetrafile.readlines()

tetradict = {}

for so in genome:
    tetradict = {}
    for t in tetras:
        tet = t.strip()
        tetradict[tet] = 0
        st = Motif.Motif(alphabet=IUPAC.unambiguous_dna)
        st.add_instance(Seq(tet, st.alphabet))
        for pos, seq in st.search_instances(so.seq):
            tetradict[seq.tostring()] += 1

#        print "Done with ", tet, tetradict[tet]
    tetralist = tetradict.items()
    tetralist.sort()
    total = 0
    numbers = []

    for i in tetralist:
        numbers.append(i[1])

    total = numpy.sum(numbers)
    average = numpy.average(numbers)
    stdev = numpy.std(numbers)

    toprint = so.id
    for j in tetralist:
        z = zscore(j[1], average, stdev)
        toprint += "\t" + '{0:.8}'.format(str(z))
    print toprint

#for k, v in tetradict.iteritems():
#    print k, tetradict[k]


tetrafile.close()
    
